use std::path::PathBuf;

use burn::{
    config::Config,
    module::Module,
    prelude::Backend,
    record::{CompactRecorder, Recorder},
    tensor::cast::ToElement,
};

use crate::{
    data::{batcher::HWDBBatcher, utils::read_classes_file_from_hdf5},
    train::train::TrainingConfig,
};

pub fn infer<B: Backend>(
    artifact_dir: &str,
    device: B::Device,
    images_dir: Vec<String>,
    h5_filepath: String,
    save_images: Option<String>,
) {
    let classes = read_classes_file_from_hdf5(&h5_filepath).expect("read classes file failed");

    let config =
        TrainingConfig::load(format!("{artifact_dir}/config.json")).expect("Failed to load config");
    let record = CompactRecorder::new()
        .load(format!("{artifact_dir}/model").into(), &device)
        .expect("Failed to load model");
    let model = config.model.init::<B>(&device).load_record(record);

    let batcher = HWDBBatcher::new(
        false,
        config.image_size,
        config.image_size,
        Some((0.9398248, 0.47693425, 269, 259)),
    );

    let save_images = save_images
        .map(|path| PathBuf::from(&path))
        .filter(|path| path.exists());
    let batch = batcher.transfer_batch(images_dir, &device, save_images);
    let outpt = model.forward(batch);
    let pred = outpt.argmax(1).flatten::<1>(0, 1);
    let label: Vec<String> = pred
        .clone()
        .iter_dim(0)
        .map(|x| x.into_scalar().to_usize())
        .map(|x| classes.get(x).expect("No class").clone())
        .collect();

    println!("Prediction: {}, Label: {:?}", pred, label);
}
